"""multi_modal_models_citation_pipeline.py"""

import asyncio
import re
import sys
import os
import json
import nest_asyncio
from jinja2 import Environment
from research_agent.core.find_statement_citation import FindStatementCitation
from research_agent.core.add_citation import AddCitation
from research_agent.core.verify_statement_citation import StatementCitationVerifier
from research_agent.core.update_reference import UpdateReference
import json
from functools import wraps
import logging




def write_to_file(text, filename):
    with open(filename, 'w') as f:
        json.dump(text, f,indent=4)


def async_retry(retries=3, delay=1):
    """
    异步重试装饰器
    Args:
        retries (int): 最大重试次数
        delay (int): 重试间隔时间(秒)
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(retries):
                try:
                    result = await func(*args, **kwargs)
                    return result
                except Exception as e:
                    last_exception = e
                    if attempt < retries - 1:
                        await asyncio.sleep(delay)
                    logging.warning(f"第 {attempt + 1} 次尝试失败: {str(e)}")
            raise last_exception
        return wrapper
    return decorator

class MultiModalCitationPipeline:
    def __init__(self, paper_draft: str, topic: str, max_retries: int = 3,logger=None):
        # 构建论文草稿
        if logger is None:
            self.logger = logging.getLogger(__name__)
            self.logger.setLevel(logging.DEBUG)
            handler = logging.StreamHandler()
            handler.setLevel(logging.DEBUG)
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            if not self.logger.handlers:
                self.logger.addHandler(handler)
        else:
            self.logger = logger
        self.paper_draft = paper_draft
        self.topic = topic
        # 初始化核心组件
        self.find_statement_citation = FindStatementCitation()
        self.add_citation = AddCitation()
        self.verifier = StatementCitationVerifier()
        self.updater = UpdateReference()
        # 配置参数
        self.max_retries = max_retries
        # self.survey_info = {"section":self.paper_draft}

    async def run_pipeline(self) -> dict:
        """完整执行引文处理流程"""
        results = {}
        
        # 步骤1: 初始自动化注释分析
        @async_retry(retries=3, delay=1)
        async def find_statements():
            return await self._find_statements()
        find_statements_results = await find_statements()
        results["initial_with_citations"] = len(find_statements_results)
        self.logger.debug(f"find statements: {len(find_statements_results)}")

        # 步骤2: 初次补充引文
        @async_retry(retries=3, delay=1)
        async def supplement_citations(statements):
            return await self._supplement_citations(statements)
        supp_results = await supplement_citations(find_statements_results)
        self.logger.debug(f"supplement citations: {len(supp_results)}")

        # 步骤3: 引文验证循环
        @async_retry(retries=3, delay=1)
        async def verify_citations(supplemented_statements):
            return await self._verify_citations(
                supplemented_statements,
                remaining_retries=self.max_retries
            )
        verified_results = await verify_citations(supp_results)
        results.update(verified_results)
        self.logger.debug(f"verified results: {verified_results}")
        filter_results = []
        for v in verified_results["supported"]:
            a_re = re.findall('<sup>',v["statement"])
            if not a_re:
                continue
            else:
                filter_results.append(v)
        verified_results["supported"] = filter_results
        # 步骤4: 最终更新论文草稿
        @async_retry(retries=3, delay=1)
        async def update_draft(supported):
            return await self._update_draft(supported)
        final_draft = await update_draft(verified_results["supported"])
        results["final_draft"] = final_draft
        self.logger.debug(f"Successfully updated {self.paper_draft[:50]} with citations")

        return results["final_draft"]

    async def _find_statements(self):
        """封装引用查找操作"""
        return await self.find_statement_citation.find_statement_citation(
            self.topic,
            self.paper_draft
        )

    async def _supplement_citations(self, statements: list):
        """封装引文补充操作"""
        return await self.add_citation.add_citations(statements)

    async def _verify_citations(self, citations: list, remaining_retries: int) -> dict:
        """递归验证流程"""
        supported, unsupported = await self.verifier.verify_statements(citations)
        results = {
            "supported": supported,
            "unsupported_count": len(unsupported),
            "retries_remaining": self.max_retries,
        }

        while unsupported and results["retries_remaining"] > 0:
            new_supplemented_citations = await self._supplement_citations(unsupported)
            new_verified_supported, new_verified_unsupported = await self.verifier.verify_statements(new_supplemented_citations)
            results["supported"] += new_verified_supported
            results["unsupported_count"] = len(new_verified_unsupported)
            results["retries_remaining"] -= 1
            if results["retries_remaining"] == 0 and new_verified_unsupported:
                # 将未支持的语句和其支持的论文引用作为字符串加入 supported
                results["supported"] += [{"statement":f"{unsupported_statement['statement']}{''.join(unsupported_statement['supported_papers'])}"}
                    for unsupported_statement in new_verified_unsupported
                ]
            unsupported = new_verified_unsupported

        return results

    async def _update_draft(self, supported: list) -> str:
        """最终草稿更新操作"""
        update_result = await self.updater.update_reference(supported, self.paper_draft)
        return update_result  


# with open("survey_info.json", "r", encoding="utf-8") as f:
#     survey_info = json.load(f)
# paper_draft = "\n------\n".join([
#     survey_info["introduction_question1"],
#     survey_info["related_work_question1"],
#     survey_info["conclusion_question1"]
# ])
# topic = "What does the technology development roadmap for multi-modal large models look like?"
# max_retries = 3
# 使用示例
# if __name__ == "__main__":
#     async def main():

#         pipeline = MultiModalCitationPipeline(paper_draft, topic, max_retries)
#         results = await pipeline.run_pipeline()

#         print("\n=== 处理结果 ===")
#         if "error" in results:
#             print(f"错误: {results['error']}")
#         else:
#             print(f"初始有引用的声明数: {results['initial_with_citations']}")
#             print(f"初始缺失引用的声明数: {results['initial_without_citations']}")
#             print(f"最终支持度验证后的声明数: {len(results['supported'])}")
#             print(f"经过 {results['retries_remaining']} 次验证循环后")
#             print("\n=== 最终草稿 ===")
#             print(results['final_draft'][:500] + "...")  # 展示部分内容

#     asyncio.run(main())
